package worker

import (
	"context"
	"gitee.com/Cookie_XiaoD/crontab/common"
	"github.com/coreos/etcd/clientv3"
)

//分布式任务锁
type JobLock struct {
	kv    clientv3.KV
	lease clientv3.Lease
	//任务名
	jobName string
	//租约ID
	leaseID clientv3.LeaseID
	//用于取消锁的自动续租
	cancelFunc context.CancelFunc
	//是否上锁成功
	isLocked bool
}

//创建初始化任务锁
func InitJobLock(jobName string, kv clientv3.KV, lease clientv3.Lease) (jobLock *JobLock) {
	jobLock = &JobLock{
		kv:      kv,
		jobName: jobName,
		lease:   lease,
	}
	return
}

//尝试获取乐观锁
//若获取不到不重试
func (jobLock *JobLock) TryLock() (err error) {
	//etcd分布式乐观锁
	//创建租约，通过租约保证锁可以过期，否者获取锁的客户端挂掉，则永远占有锁
	//占用锁的过程中需要自动续租，保证锁独占
	//锁的路径为 /cron/locks/任务名
	//通过etcd的txn事务抢锁
	var (
		leaseResp  *clientv3.LeaseGrantResponse
		cancelCtx  context.Context
		cancelFunc context.CancelFunc
		leaseID    clientv3.LeaseID
		keepChan   <-chan *clientv3.LeaseKeepAliveResponse
		txn        clientv3.Txn
		lockKey    string
		txnResp    *clientv3.TxnResponse
	)
	//申请租约 5秒
	if leaseResp, err = jobLock.lease.Grant(context.TODO(), 5); err != nil {
		return
	}

	//取消上下文，用于取消自动续租
	cancelCtx, cancelFunc = context.WithCancel(context.TODO())
	//租约ID，用于续租和停止
	leaseID = leaseResp.ID

	//自动续租
	if keepChan, err = jobLock.lease.KeepAlive(cancelCtx, leaseID); err != nil {
		//自动续租失败，取消自动续租并释放租约
		goto FAIL
	}

	//通过协程处理自动续租应答
	go func() {
		//自动续租成功，处理自动续租应答
		var kaResp *clientv3.LeaseKeepAliveResponse

		for {
			select {
			case kaResp = <-keepChan:
				if kaResp == nil { //自动续租失败
					goto END
				}
			}
		}
	END:
	}()

	//通过txn事务获取锁
	lockKey = common.JOB_LOCK_DIR + jobLock.jobName
	txn = jobLock.kv.Txn(context.TODO())
	txn.If(clientv3.Compare(clientv3.CreateRevision(lockKey), "=", 0)).
		Then(clientv3.OpPut(lockKey, "", clientv3.WithLease(leaseID))).
		Else(clientv3.OpGet(lockKey))

	if txnResp, err = txn.Commit(); err != nil {
		goto FAIL //发生错误
	}

	//事务结果失败
	if !txnResp.Succeeded {
		err = common.ERR_LOCK_ALREADY_REQUIRED
		return
	}

	//事务执行成功，获取锁成功
	jobLock.cancelFunc=cancelFunc
	jobLock.isLocked=true
	jobLock.leaseID=leaseID

	return

FAIL:
	cancelFunc()
	jobLock.lease.Revoke(context.TODO(), leaseID)
	return

}

//释放锁
func (jobLock *JobLock) Unlock() {
	if jobLock.isLocked {
		//释放锁就是取消续租，释放租约
		jobLock.cancelFunc()
		jobLock.lease.Revoke(context.TODO(), jobLock.leaseID)
	}
}
